import numpy as np
import matplotlib.pyplot as plt
from sklearn.svm import SVC
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split

rnd_seed = 2

# load
data = np.loadtxt(r'../data/xigua.txt', delimiter=',')
x = data[:, :-1]
y = data[:, -1]

# scale
scaler = StandardScaler()
x = scaler.fit_transform(x)

# split
x_train, x_test, y_train, y_test = train_test_split(x, y, train_size=0.8, random_state=rnd_seed)

# train
clf = SVC(gamma=1)
clf.fit(x_train, y_train)
print(f'Training score = {clf.score(x_train, y_train)}')

# test
print(f'Testing score = {clf.score(x_test, y_test)}')

xx, yy = np.mgrid[x[:, 0].min():x[:, 0].max():100j, x[:, 1].min():x[:, 1].max():100j]
xxyy = np.c_[xx.ravel(), yy.ravel()]
zzf = clf.predict(xxyy).reshape(xx.shape)
zz = clf.decision_function(xxyy).reshape(xx.shape)

plt.figure(figsize=[16, 8])
spr = 1  # subplot row
spc = 2  # subplot column
spn = 0  # subplot number

spn += 1
plt.subplot(spr, spc, spn)
bi_cmap = plt.cm.get_cmap('rainbow', 2)
plt.scatter(x_train[:, 0], x_train[:, 1], s=100, marker='x', zorder=100, c=y_train, cmap=bi_cmap)
plt.contour(xx, yy, zz, cmap=plt.cm.Paired, levels=20, zorder=0)

spn += 1
plt.subplot(spr, spc, spn)
plt.scatter(x_test[:, 0], x_test[:, 1], s=100, marker='x', zorder=100, c=y_test, cmap=bi_cmap)
plt.contourf(xx, yy, zzf, cmap=plt.cm.Paired, zorder=0)

plt.show()
